package cn.aicnn.chatssespringboot.service.Impl;

import cn.aicnn.chatssespringboot.service.BaiduAIPainting;
import okhttp3.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@Service
public class BaiduAIPaintingImpl implements BaiduAIPainting {

    private static final String API_KEY = "BDh4GNJ9lNYXNkEQSh6N7zCV";
    private static final String SECRET_KEY = "7H3YhCeZ0wfW6uZTUtiigjMkDNwKOuFA";

//    public static void main(String[] args) throws IOException, InterruptedException {
//        String taskId = submitRequest("帮我画一副八仙过海图");
//        System.out.println("Task ID: " + taskId);
//
//        List<String> imgUrls = queryResult(taskId);
//        String savedImagePath = downloadAndSaveImage(imgUrls, "C:/Users/z's'j/Desktop/images");
//        System.out.println("Image saved at: " + savedImagePath);
//    }

    public String submitRequest(String promot) throws IOException {
        String token = getAccessToken();

        OkHttpClient client = new OkHttpClient();
        ObjectMapper objectMapper = new ObjectMapper();

        String url = "https://aip.baidubce.com/rpc/2.0/ernievilg/v1/txt2img?access_token=" + token;

        Map<String, Object> payloadMap = new HashMap<>();
        payloadMap.put("text", promot);
        payloadMap.put("resolution", "1024*1024");
        payloadMap.put("style", "写实风格");
        payloadMap.put("num", 1);

        String payload = objectMapper.writeValueAsString(payloadMap);

        RequestBody body = RequestBody.create(payload, MediaType.parse("application/json"));

        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .build();

        try (Response response = client.newCall(request).execute()) {
            String responseBody = response.body().string();
            System.out.println("Response Body: " + responseBody); // 打印响应体以便调试

            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response + " with body " + responseBody);
            }

            Map<String, Object> result = objectMapper.readValue(responseBody, HashMap.class);
            Map<String, Object> data = (Map<String, Object>) result.get("data");

            if (data == null) {
                throw new IllegalStateException("Response data is null: " + result);
            }

            Object taskIdObject = data.get("taskId");
            if (taskIdObject instanceof Integer) {
                return String.valueOf(taskIdObject);
            } else if (taskIdObject instanceof String) {
                return (String) taskIdObject;
            } else {
                throw new IllegalStateException("Unexpected taskId type: " + (taskIdObject != null ? taskIdObject.getClass().getName() : "null"));
            }
        }
    }

    public List<String> queryResult(String taskId) throws IOException, InterruptedException {
        String token = getAccessToken();
        OkHttpClient client = new OkHttpClient();
        ObjectMapper objectMapper = new ObjectMapper();
        String url = "https://aip.baidubce.com/rpc/2.0/ernievilg/v1/getImg?access_token=" + token;

        while (true) {
            Map<String, Object> payloadMap = new HashMap<>();
            payloadMap.put("taskId", taskId);
            String payload = objectMapper.writeValueAsString(payloadMap);
            RequestBody body = RequestBody.create(payload, MediaType.parse("application/json"));
            Request request = new Request.Builder()
                    .url(url)
                    .post(body)
                    .addHeader("Content-Type", "application/json")
                    .addHeader("Accept", "application/json")
                    .build();

            try (Response response = client.newCall(request).execute()) {
                String responseBody = response.body().string();
                System.out.println("Response Body: " + responseBody); // 打印响应体以便调试

                if (!response.isSuccessful()) {
                    throw new IOException("Unexpected code " + response + " with body " + responseBody);
                }

                Map<String, Object> result = objectMapper.readValue(responseBody, HashMap.class);
                Map<String, Object> data = (Map<String, Object>) result.get("data");

                if (data == null) {
                    throw new IllegalStateException("Response data is null: " + result);
                }

                List<Map<String, String>> imgUrls = (List<Map<String, String>>) data.get("imgUrls");
                if (imgUrls != null && !imgUrls.isEmpty()) {
                    return imgUrls.stream()
                            .map(imgUrlMap -> imgUrlMap.get("image"))
                            .toList();
                }

                int status = (int) data.get("status");
                if (status == 0) {
                    // 如果任务未完成，等待一段时间再查询
                    System.out.println("Task is not complete, waiting for 10 seconds...");
                    Thread.sleep(10000);
                } else {
                    throw new IllegalStateException("Unexpected task status: " + status);
                }
            }
        }
    }

    public String downloadAndSaveImage(List<String> imgUrls, String folderPath) throws IOException {
        if (imgUrls == null || imgUrls.isEmpty()) {
            throw new IllegalArgumentException("Image URLs list is empty or null.");
        }

        OkHttpClient client = new OkHttpClient();
        File folder = new File(folderPath);
        if (!folder.exists()) {
            folder.mkdirs();
        }

        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMddHHmmssSSS");
        String imgUrl = imgUrls.get(0); // Assuming there's only one URL in the list

        Request request = new Request.Builder()
                .url(imgUrl)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response);
            }

            InputStream inputStream = response.body().byteStream();
            String timestamp = sdf.format(new Date());
            File imageFile = new File(folder, "image_" + timestamp + ".jpg");

            try (FileOutputStream outputStream = new FileOutputStream(imageFile)) {
                byte[] buffer = new byte[2048];
                int length;
                while ((length = inputStream.read(buffer)) != -1) {
                    outputStream.write(buffer, 0, length);
                }
            }

            System.out.println("Image saved to " + imageFile.getAbsolutePath());
            return imageFile.getAbsolutePath(); // Return the saved image path
        }
    }

    public static String getAccessToken() throws IOException {
        OkHttpClient client = new OkHttpClient();

        String url = "https://aip.baidubce.com/oauth/2.0/token";
        HttpUrl.Builder httpBuilder = HttpUrl.parse(url).newBuilder();
        httpBuilder.addQueryParameter("grant_type", "client_credentials");
        httpBuilder.addQueryParameter("client_id", API_KEY);
        httpBuilder.addQueryParameter("client_secret", SECRET_KEY);

        Request request = new Request.Builder()
                .url(httpBuilder.build())
                .post(RequestBody.create(new byte[0]))
                .build();

        try (Response response = client.newCall(request).execute()) {
            String responseBody = response.body().string();
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response + " with body " + responseBody);
            }

            Map<String, Object> result = new ObjectMapper().readValue(responseBody, HashMap.class);
            return (String) result.get("access_token");
        }
    }
}
